package vn.bum.game.web.base;

import java.io.ByteArrayOutputStream;
import java.io.CharArrayWriter;
import java.io.IOException;
import java.io.PrintWriter;

import javax.servlet.ServletOutputStream;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;

public class BufferedResponseWrapper extends HttpServletResponseWrapper {
	private ByteArrayOutputStream byteBuffer;
	private CharArrayWriter charBuffer;
	private ServletOutputStream outputStream;
	private PrintWriter writer;

	public BufferedResponseWrapper(ServletResponse response) {
		super((HttpServletResponse) response);
	}

	/**
	 * Returns a binary output stream
	 */
	public ServletOutputStream getOutputStream() throws IOException {
		if (writer != null) {
			String errmsg = "getWriter() was already called";
			throw new IllegalStateException(errmsg);
		}

		if (outputStream == null) {

			// Create a subclass of ServletOutputStream
			// that delegates everything to a
			// ByteArrayOutputStream
			byteBuffer = new ByteArrayOutputStream();
			outputStream = new ServletOutputStream() {
				public void write(int c) throws IOException {
					byteBuffer.write(c);
				}
			};
		}

		return outputStream;
	}

	/**
	 * Returns a character writer
	 */
	public PrintWriter getWriter() throws IOException {
		if (outputStream != null) {
			String errmsg = "getOutputStream() was already called";
			throw new IllegalStateException(errmsg);
		}

		if (writer == null) {
			charBuffer = new CharArrayWriter();
			writer = new PrintWriter(charBuffer);
		}

		return writer;
	}

	/**
	 * Returns the contents of the buffer as a string
	 */
	public String getBufferAsString() throws IOException {
		String buffer = null;

		if (charBuffer != null) {
			writer.flush();
			charBuffer.flush();
			buffer = charBuffer.toString();
		} else if (byteBuffer != null) {
			outputStream.flush();
			byteBuffer.flush();
			buffer = byteBuffer.toString();
		} else
			buffer = "";

		return buffer;
	}
}
